#include <cstdio>
#include <cstdlib>
#include <cuda_runtime.h>
#include <vector>

#define ALIGNMENT 8
#define POOL_SIZE (1 << 25)  // 16MB 内存池
#define MIN_BLOCK_SIZE (sizeof(size_t) * 3) // 最小块大小（头部+尾部+最小数据）

// 简化的内存池管理器
struct SimplePoolManager {
    char* pool_start;     // 内存池起始地址
    size_t pool_size;     // 内存池总大小
    void* free_list;      // 自由链表头指针
    int lock;             // 自旋锁
};

// 设备端全局指针
__device__ SimplePoolManager* device_pool_manager = nullptr;

// 自旋锁获取
__device__ void lock_acquire(int* lock) {
    while (atomicCAS(lock, 0, 1) != 0);
}

// 自旋锁释放
__device__ void lock_release(int* lock) {
    atomicExch(lock, 0);
}

// 对齐计算
__device__ size_t align_up(size_t size, size_t alignment) {
    return (size + alignment - 1) & ~(alignment - 1);
}

// 获取块大小（清除标志位）
__device__ size_t get_block_size(void* block) {
    return *((size_t*)block) & ~1;
}

// 获取块状态
__device__ bool is_block_free(void* block) {
    return *((size_t*)block) & 1;
}

// 设置块头部和尾部
__device__ void set_block_header(void* block, size_t size, bool is_free) {
    size_t header_val = size | (is_free ? 1 : 0);
    *((size_t*)block) = header_val; // 头部
    *((size_t*)((char*)block + size - sizeof(size_t))) = header_val; // 尾部
}

// 初始化空闲块（设置头部、尾部和链表指针）
__device__ void init_free_block(void* block, size_t size, void* next) {
    set_block_header(block, size, true);
    void** next_ptr = (void**)((char*)block + sizeof(size_t));
    *next_ptr = next;
}

// 合并相邻空闲块
__device__ void coalesce_free_blocks(void* block) {
    SimplePoolManager* mgr = device_pool_manager;
    size_t block_size = get_block_size(block);
    char* block_end = (char*)block + block_size;

    // 检查后一块是否空闲
    if (block_end < mgr->pool_start + mgr->pool_size) {
        void* next_block = block_end;
        if (is_block_free(next_block)) {
            size_t next_block_size = get_block_size(next_block);
            // 从自由链表移除后一块
            void** prev_ptr = &mgr->free_list;
            void* curr = mgr->free_list;
            while (curr) {
                void** next_ptr = (void**)((char*)curr + sizeof(size_t));
                if (curr == next_block) {
                    *prev_ptr = *next_ptr; // 从链表移除
                    break;
                }
                prev_ptr = next_ptr;
                curr = *next_ptr;
            }
            // 合并块
            block_size += next_block_size;
            set_block_header(block, block_size, true);
        }
    }

    // 检查前一块是否空闲
    if (block > mgr->pool_start) {
        void* prev_block = (char*)block - sizeof(size_t);
        size_t prev_block_size = *((size_t*)prev_block) & ~1;
        prev_block = (char*)block - prev_block_size;
        if (is_block_free(prev_block)) {
            block = prev_block;
            block_size += prev_block_size;
            // 从自由链表移除前一块
            void** prev_ptr = &mgr->free_list;
            void* curr = mgr->free_list;
            while (curr) {
                void** next_ptr = (void**)((char*)curr + sizeof(size_t));
                if (curr == prev_block) {
                    *prev_ptr = *next_ptr; // 从链表移除
                    break;
                }
                prev_ptr = next_ptr;
                curr = *next_ptr;
            }
            // 更新合并块
            set_block_header(block, block_size, true);
        }
    }

    // 将合并后的块插入链表头部
    init_free_block(block, block_size, mgr->free_list);
    mgr->free_list = block;
}

// 简化的内存分配
__device__ void* simple_malloc(size_t size) {
    if (size == 0 || device_pool_manager == nullptr) 
        return nullptr;

    size_t aligned_size = align_up(size, ALIGNMENT);
    size_t required_size = aligned_size + sizeof(size_t) * 2; // 头部+尾部
    if (required_size < MIN_BLOCK_SIZE) required_size = MIN_BLOCK_SIZE;

    SimplePoolManager* mgr = device_pool_manager;
    lock_acquire(&mgr->lock);

    void** prev_ptr = &mgr->free_list;
    void* curr = mgr->free_list;
    void* found_block = nullptr;
    size_t found_size = 0;
    size_t remaining_size = 0;

    // 首次适应搜索
    while (curr) {
        void** next_ptr = (void**)((char*)curr + sizeof(size_t));
        size_t block_size = get_block_size(curr);
        if (block_size >= required_size) {
            found_block = curr;
            found_size = block_size;
            // 从链表移除
            *prev_ptr = *next_ptr;
            break;
        }
        prev_ptr = next_ptr;
        curr = *next_ptr;
    }

    if (found_block) {
        // 分割块（剩余部分至少为最小块）
        remaining_size = found_size - required_size;
        if (remaining_size >= MIN_BLOCK_SIZE) {
            // 设置分配块
            set_block_header(found_block, required_size, false);
            // 设置剩余空闲块
            void* new_free_block = (char*)found_block + required_size;
            init_free_block(new_free_block, remaining_size, mgr->free_list);
            mgr->free_list = new_free_block;
        } else {
            // 不分割，整块分配
            set_block_header(found_block, found_size, false);
        }
        // 返回用户数据地址（头部之后）
        found_block = (char*)found_block + sizeof(size_t);
    }

    lock_release(&mgr->lock);
    return found_block;
}

// 内存释放
__device__ void simple_free(void* ptr) {
    if (!ptr || !device_pool_manager) 
        return;

    // 获取块起始地址（头部）
    void* block = (char*)ptr - sizeof(size_t);
    SimplePoolManager* mgr = device_pool_manager;
    lock_acquire(&mgr->lock);
    // 标记为空闲并合并
    coalesce_free_blocks(block);
    lock_release(&mgr->lock);
}

// 重置内存池
__device__ void reset_pool() {
    if (device_pool_manager) {
        SimplePoolManager* mgr = device_pool_manager;
        lock_acquire(&mgr->lock);
        // 整个池作为单个空闲块
        init_free_block(mgr->pool_start, mgr->pool_size, nullptr);
        mgr->free_list = mgr->pool_start;
        lock_release(&mgr->lock);
    }
}

// 重置内存池的核函数
__global__ void reset_pool_kernel() {
    reset_pool();
}

// 初始化全局内存池
SimplePoolManager* create_simple_pool(size_t pool_size) {
    char* d_pool;
    cudaMalloc(&d_pool, pool_size);
    
    SimplePoolManager* d_manager;
    cudaMalloc(&d_manager, sizeof(SimplePoolManager));
    
    SimplePoolManager h_manager;
    h_manager.pool_start = d_pool;
    h_manager.pool_size = pool_size;
    h_manager.free_list = nullptr;
    h_manager.lock = 0;
    
    cudaMemcpy(d_manager, &h_manager, sizeof(SimplePoolManager), cudaMemcpyHostToDevice);
    cudaMemcpyToSymbol(device_pool_manager, &d_manager, sizeof(SimplePoolManager*));
    
    // 初始化设备端自由链表
    reset_pool_kernel<<<1, 1>>>();
    cudaDeviceSynchronize();
    
    return d_manager;
}

// 销毁全局内存池
void destroy_simple_pool(SimplePoolManager* manager) {
    char* pool_start;
    cudaMemcpy(&pool_start, &manager->pool_start, sizeof(char*), cudaMemcpyDeviceToHost);
    
    cudaFree(pool_start);
    cudaFree(manager);
    
    SimplePoolManager* null_ptr = nullptr;
    cudaMemcpyToSymbol(device_pool_manager, &null_ptr, sizeof(SimplePoolManager*));
}

// 打印内存池状态
void print_pool_status(SimplePoolManager* manager) {
    SimplePoolManager h_manager;
    cudaMemcpy(&h_manager, manager, sizeof(SimplePoolManager), cudaMemcpyDeviceToHost);
    
    size_t total_free = 0;
    void* curr = h_manager.free_list;
    while (curr) {
        size_t block_size;
        cudaMemcpy(&block_size, curr, sizeof(size_t), cudaMemcpyDeviceToHost);
        block_size &= ~1;
        total_free += block_size;
        void* next;
        cudaMemcpy(&next, (char*)curr + sizeof(size_t), sizeof(void*), cudaMemcpyDeviceToHost);
        curr = next;
    }
    
    size_t used = h_manager.pool_size - total_free;
    
    printf("FreeList Memory Pool Status:\n");
    printf("  Total Size: %zu bytes\n", h_manager.pool_size);
    printf("  Used: %zu bytes (%.1f%%)\n", used, (used * 100.0) / h_manager.pool_size);
    printf("  Free: %zu bytes (%.1f%%)\n", total_free, (total_free * 100.0) / h_manager.pool_size);
}


// 测试核函数
__global__ void test_kernel(int* results, size_t num_allocations) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (tid < num_allocations) {
        // 分配不同大小的内存
        size_t alloc_size = (tid % 4 + 1) * 64;  // 64/128/192/256 bytes
        
        // 分配内存
        char* ptr = (char*)simple_malloc(alloc_size);
        
        if (ptr) {
            // 使用内存
            for (size_t i = 0; i < alloc_size; i++) {
                ptr[i] = (tid + i) % 256;
            }
            
            // 验证写入
            bool valid = true;
            for (size_t i = 0; i < alloc_size; i++) {
                if (ptr[i] != (char)((tid + i) % 256)) {
                    valid = false;
                    break;
                }
            }
            
            results[tid] = valid ? 1 : -1;
        } else {
            results[tid] = 0;  // 分配失败
        }
    }
}



int main() {
    const size_t pool_size = POOL_SIZE;
    const size_t num_allocations = 100000;
    
    // 创建全局内存池
    SimplePoolManager* pool = create_simple_pool(pool_size);
    printf("Created simple memory pool (%zu bytes)\n", pool_size);
    
    // 分配设备内存存储结果
    int* d_results;
    cudaMalloc(&d_results, num_allocations * sizeof(int));
    
    // 第一次运行测试
    int block_size = 256;
    int grid_size = (num_allocations + block_size - 1) / block_size;
    
    test_kernel<<<grid_size, block_size>>>(d_results, num_allocations);
    cudaDeviceSynchronize();
    printf("First kernel completed\n");
    
    // 检查结果
    std::vector<int> h_results(num_allocations);
    cudaMemcpy(h_results.data(), d_results, 
               num_allocations * sizeof(int), cudaMemcpyDeviceToHost);
    
    int success = 0, fail = 0, alloc_fail = 0;
    for (int r : h_results) {
        if (r == 1) success++;
        else if (r == -1) fail++;
        else alloc_fail++;
    }
    
    printf("First run results:\n");
    printf("  Success: %d\n", success);
    printf("  Fail: %d\n", fail);
    printf("  Allocation failures: %d\n", alloc_fail);
    print_pool_status(pool);
    
    // 重置内存池
    reset_pool_kernel<<<1, 1>>>();
    cudaDeviceSynchronize();
    printf("\nPool reset\n");
    print_pool_status(pool);
    
    // 第二次运行测试（不同block配置）
    int block_size2 = 128;
    int grid_size2 = (num_allocations + block_size2 - 1) / block_size2;
    
    test_kernel<<<grid_size2, block_size2>>>(d_results, num_allocations);
    cudaDeviceSynchronize();
    printf("\nSecond kernel completed (different block config)\n");
    
    // 检查结果
    cudaMemcpy(h_results.data(), d_results, 
               num_allocations * sizeof(int), cudaMemcpyDeviceToHost);
    
    success = fail = alloc_fail = 0;
    for (int r : h_results) {
        if (r == 1) success++;
        else if (r == -1) fail++;
        else alloc_fail++;
    }
    
    printf("Second run results:\n");
    printf("  Success: %d\n", success);
    printf("  Fail: %d\n", fail);
    printf("  Allocation failures: %d\n", alloc_fail);
    print_pool_status(pool);
    
    // 清理
    cudaFree(d_results);
    destroy_simple_pool(pool);
    
    printf("\nMemory pool destroyed\n");
    return 0;
}